#pragma once
#undef slots
#include <torch/torch.h>
#include <torch/script.h>
#include "Vocab.h"
#define slots Q_SLOTS
#include <string>
#include <vector>

class chEnModel {
public:
    chEnModel() = default;
    ~chEnModel() = default;
    explicit chEnModel(const std::string& model_path, const Vocab& src_vocab, const Vocab& tar_vocab);

    std::string predict(const std::string& src_sentence, int num_steps = 64);

private:
    torch::jit::Module m_model;
    const Vocab& m_src_vocab;
    const Vocab& m_tar_vocab;
    torch::Device m_device = torch::cuda::is_available() ? torch::kCUDA : torch::kCPU;

    torch::Tensor preprocess_input(const std::string& src_sentence, int num_steps);
};